import random
def assign_cluster(x, centers):
    min_dist_sq = float('inf')
    best_cluster_idx = 0

    # 计算样本到每个聚类中心的欧氏距离平方
    for idx, center in enumerate(centers):
        dist_sq = sum((xi - ci) ** 2 for xi, ci in zip(x, center))
        if dist_sq < min_dist_sq:
            min_dist_sq = dist_sq
            best_cluster_idx = idx

    return best_cluster_idx

def Kmeans(data, k, epsilon=1e-3, iteration=100):
    # 1. 输入参数校验
    if not data:
        raise ValueError("输入数据不能为空")

    n_samples = len(data)
    if k < 1 or k > n_samples:
        raise ValueError(f"k值必须满足 1 ≤ k ≤ 样本数（当前样本数：{n_samples}）")

    n_features = len(data[0])
    for sample in data:
        if len(sample) != n_features:
            raise ValueError("所有样本必须具有相同的维度")

    # 2. 初始化聚类中心
    init_indices = random.sample(range(n_samples), k)
    final_centers = [data[i].copy() for i in init_indices]

    # 3. 迭代优化聚类中心
    for iter_cnt in range(iteration):

        clusters = [[] for _ in range(k)]
        labels = []

        for sample_idx, sample in enumerate(data):

            cluster_idx = assign_cluster(sample, final_centers)
            clusters[cluster_idx].append(sample_idx)
            labels.append(cluster_idx)

        old_centers = [center.copy() for center in final_centers]

        for cluster_idx in range(k):
            cluster_samples = clusters[cluster_idx]

            if not cluster_samples:
                final_centers[cluster_idx] = data[random.randint(0, n_samples - 1)].copy()
                continue

            new_center = []
            for dim in range(n_features):
                dim_sum = sum(data[s_idx][dim] for s_idx in cluster_samples)
                new_center.append(dim_sum / len(cluster_samples))
            final_centers[cluster_idx] = new_center

        total_center_dist = 0.0

        for old_c, new_c in zip(old_centers, final_centers):
            dist = sum((o - n) ** 2 for o, n in zip(old_c, new_c)) ** 0.5
            total_center_dist += dist

        if total_center_dist < epsilon:
            print(f"迭代 {iter_cnt + 1} 次后收敛（中心总变化量：{total_center_dist:.6f} < ε={epsilon}）")
            break

    else:
        print(f"已达到最大迭代次数 {iteration}，未完全收敛（最终中心总变化量：{total_center_dist:.6f}）")

    return labels, final_centers